import re
import os
from typing import Dict, Any, List


def count_node_occurrences(tree: Dict[str, Any]) -> Dict[str, int]:
    """统计树中每个节点的出现次数"""
    counter = {}
    
    def traverse(node):
        # 提取节点名称（移除前缀标记）
        name = node['name']
        if name.startswith('repeat## '):
            name = name[9:]  # 移除'repeat## '前缀
        elif name.startswith('** '):
            name = name[3:]  # 移除'** '前缀
            
        # 增加计数
        counter[name] = counter.get(name, 0) + 1
        
        # 遍历子节点
        for child in node.get('children', []):
            traverse(child)
    
    traverse(tree)
    return counter


def print_top_calls(tree: Dict[str, Any], root_name: str, n: int = 10):
    """打印调用次数最多的前n个函数及其调用次数，并返回名称列表"""
    counter= count_node_occurrences(tree)
    # 按出现次数降序排序
    sorted_items = sorted(counter.items(), key=lambda x: x[1], reverse=True)
    
    # 提取前n个 
    top_n = sorted_items[:n]
    
    # 格式化输出
    names = [item[0] for item in top_n]  # 函数名称
    counts = [item[1] for item in top_n]  # 调用次数
    
    print(f"\n{root_name}调用频率最高的前{n}个函数:")
    print(names)  # 打印函数名称
    print(f"{root_name}调用次数:")
    print(counts)  # 打印调用次数
    
    # 返回名称列表，以便写入文件
    return #names, counts  # 返回函数名称和调用次数



def extract_call_chains(tree: Dict[str, Any]) -> List[str]:
    """Extract all unique function call chains from a tree structure"""
    chains = set()  # Using a set to automatically eliminate duplicates
    
    def traverse(node, current_chain=""):
        # Extract node name (remove any prefixes)
        name = node['name']
        if name.startswith('repeat## '):
            name = name[9:]
        elif name.startswith('** '):
            name = name[3:]
        
        # Update current chain
        if current_chain:
            new_chain = current_chain + "-->" + name
        else:
            new_chain = name
        
        # If leaf node, add to chains
        if not node['children']:
            chains.add(new_chain)  # Using add() for set
        else:
            # Continue traversing for each child
            for child in node['children']:
                traverse(child, new_chain)
    
    traverse(tree)
    return sorted(list(chains))  # Convert to sorted list before returning

def print_call_chains(tree: Dict[str, Any], output_file=None):
    """Print all unique function call chains from a tree structure"""
    chains = extract_call_chains(tree)
    
    # Write to file if specified
    if output_file:
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(f"\nFunction call chains from the {tree['name']} subroutine:\n")
            # f.write(f"\nFound {len(chains)} unique call chains.\n")
            for chain in chains:
                f.write(f"\nFunction call chains:\n")
                f.write(f"{chain}\n")
    
    return chains

def extract_computation_steps(file_path, start_line, end_line):
    """
    Extract computation statements from a subroutine's source code,
    including relevant comments preceding these statements.
    
    Args:
        file_path: Path to the source file
        start_line: Starting line of the subroutine
        end_line: Ending line of the subroutine
        
    Returns:
        List of computation statements with associated comments
    """
    with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
        lines = f.readlines()[start_line:end_line+1]
    
    computation_steps = []
    continuation_line = ""
    
    # Control flow keywords to exclude
    control_keywords = [
        'if', 'then', 'else', 'elseif', 'endif', 
        'do', 'enddo', 'while', 'case', 
        'select', 'endselect', 'where', 'elsewhere',
        'forall', 'exit', 'cycle', 'continue',
        'goto', 'return'
    ]
    
    # Track comments that might be relevant to the next computation statement
    pending_comments = []
    
    for line in lines:
        original_line = line.strip()
        
        # Skip the subroutine declaration and end lines
        if original_line.lower().startswith(('subroutine ', 'function ', 'end subroutine', 'end function')):
            pending_comments = []  # Reset comments at subroutine boundaries
            continue
            
        # Handle comment lines
        if line.strip().startswith('!'):
            # This is a full-line comment - save it for the next computation
            comment_text = line.strip()[1:].strip()
            if comment_text:  # Only keep non-empty comments
                pending_comments.append("! " + comment_text)
            continue
            
        # Handle inline comments
        comment_text = ""
        if '!' in line:
            parts = line.split('!', 1)
            line = parts[0]
            comment_text = parts[1].strip()
        
        # Skip empty lines
        line = line.strip()
        if not line:
            continue
            
        # Handle line continuations
        if continuation_line:
            line = continuation_line + line
            continuation_line = ""
            
        if line.endswith('&'):
            continuation_line = line[:-1]
            continue
            
        # Check if it's a control flow statement
        is_control_flow = False
        line_lower = line.lower()
        for keyword in control_keywords:
            if re.match(rf'\b{keyword}\b', line_lower):
                is_control_flow = True
                break
                
        # Include assignment statements and function calls without control flow
        if not is_control_flow and ('=' in line or '(' in line and ')' in line):
            # Clean up the statement for better readability
            clean_line = re.sub(r'\s+', ' ', line).strip()
            
            # Add inline comment if it exists
            if comment_text:
                clean_line += "  ! " + comment_text
            
            # Add any pending comments before this computation line
            if pending_comments:
                computation_steps.extend(pending_comments)
                pending_comments = []  # Reset after using
            
            computation_steps.append(clean_line)
        else:
            # For non-computation lines, discard their pending comments
            # unless they look like they might be important documentation
            if pending_comments and any(("param" in c.lower() or "input" in c.lower() or 
                                       "output" in c.lower() or "return" in c.lower() or
                                       "purpose" in c.lower() or "description" in c.lower())
                                      for c in pending_comments):
                computation_steps.extend(pending_comments)
            pending_comments = []
    
    return computation_steps

def print_all_non_leaf_nodes_steps_enhanced(tree, relationships, src_directories, 
                                          output_file=None, recursive=True):
    """
    Print the processing steps of all non-leaf nodes in a tree with actual computation steps.
    
    Args:
        tree: The tree structure to traverse
        relationships: Dictionary of function call relationships
        src_directories: List of source directories to search
        output_file: File path to save the output
        recursive: Whether to search subdirectories
    """
    # Get all non-leaf nodes in the tree
    non_leaf_nodes = set()
    visited = set()
    
    def collect_non_leaf_nodes(node):
        # Extract clean node name
        name = node['name']
        if name.startswith('repeat## '):
            name = name[9:]
        elif name.startswith('** '):
            name = name[3:]
            
        # Skip if already processed
        if name in visited:
            return
        visited.add(name)
        
        # If it has children, add to non-leaf nodes
        if node['children']:
            non_leaf_nodes.add(name)
        
        # Process all children
        for child in node['children']:
            collect_non_leaf_nodes(child)
    
    # Collect all non-leaf nodes
    collect_non_leaf_nodes(tree)
    
    # Locate all subroutines in source files
    # print(f"Locating subroutines in source directories...")
    subroutine_locations = locate_subroutines_in_files(src_directories, recursive)
    print(f"Found {len(subroutine_locations)} subroutines in source files.")
    
    # Initialize output
    output = []
    
    # Process each non-leaf node
    for node_name in sorted(non_leaf_nodes):
        # Get the direct children for this node
        calls = relationships.get(node_name, [])
        
        steps = [f"Processing steps of {node_name}:"]
        
        # Extract computation steps if the subroutine location is known
        if node_name in subroutine_locations:
            loc = subroutine_locations[node_name]
            steps.append(f"\n// Source code from {loc['file_path']}:")
            
            computation_steps = extract_computation_steps(
                loc['file_path'], loc['start_line'], loc['end_line']
            )
            
            if computation_steps:
                steps.append("// Key computation statements:")
                steps.extend(["  " + step for step in computation_steps])
            else:
                steps.append("// No computation statements found.")
        else:
            steps.append(f"// Source code not found for this subroutine.")
        
        # List all direct subfunctions with their arguments
        if calls:
            steps.append(f"\n// Function calls of {node_name}:")
            for i, call in enumerate(calls, 1):
                func_name = call['func']
                args = call.get('args', '')
                if args:
                    steps.append(f"{i}. {func_name}({args})")
                else:
                    steps.append(f"{i}. {func_name}()")
        else:
            steps.append("\n// This function doesn't call any subfunctions.")
        
        output.append("\n".join(steps))
        output.append("\n" + "-"*60 + "\n")  # Add separator
    
    # Output results
    result = "\n".join(output)
    
    # Print to console (summary only)
    print(f"Processed {len(non_leaf_nodes)} non-leaf nodes.")
    
    # Save to file if specified
    if output_file:
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(f"Processing steps with source code analysis of all non-leaf nodes:\n\n")
            f.write(result)
            print(f"Results written to {output_file}")

def locate_subroutines_in_files(src_directories, recursive=True):
    """
    Scan source directories and build a mapping of subroutine names to their file locations and line ranges.
    
    Args:
        src_directories: List of source directories to scan
        recursive: Whether to search recursively
        
    Returns:
        Dictionary mapping subroutine names to {file_path, start_line, end_line}
    """
    subroutine_locations = {}
    
    for directory in src_directories:
        for root, _, files in os.walk(directory) if recursive else [(directory, None, os.listdir(directory))]:
            for file in files:
                if file.lower().endswith(('.f90', '.f')) or file.endswith(('.F90', '.F')):
                    file_path = os.path.join(root, file)
                    with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                        lines = f.readlines()
                        
                    current_subroutine = None
                    start_line = 0
                    
                    for i, line in enumerate(lines):
                        line_lower = line.lower().strip()
                        
                        # Check for program declaration
                        if line_lower.startswith('program '):
                            program_name = line_lower.split('program', 1)[1].strip()
                            current_subroutine = program_name
                            start_line = i
                        # Check for subroutine or function declaration
                        elif (line_lower.startswith('subroutine ') or 
                            line_lower.startswith('function ') or
                            re.match(r'^\s*(?:pure|elemental|recursive)?\s+(?:subroutine|function)', line_lower)):
                            
                            # Extract name from declaration line
                            if 'subroutine' in line_lower:
                                name_parts = line_lower.split('subroutine', 1)[1].strip().split('(', 1)
                            else:
                                name_parts = line_lower.split('function', 1)[1].strip().split('(', 1)
                                
                            subroutine_name = name_parts[0].strip()
                            current_subroutine = subroutine_name
                            start_line = i
                            
                        # Check for end of subroutine or program
                        elif current_subroutine and (line_lower.startswith('end subroutine') or 
                                                    line_lower.startswith('end function') or
                                                    line_lower.startswith('end program')):
                            subroutine_locations[current_subroutine] = {
                                'file_path': file_path,
                                'start_line': start_line,
                                'end_line': i
                            }
                            current_subroutine = None
    
    return subroutine_locations



def generate_relationships(chunks: List[Dict], aliases: Dict[str, str]) -> Dict[str, List[Dict]]:
    """生成调用关系字典，包含参数"""
    relationships = {
        chunk['name']: chunk['calls']
        for chunk in chunks
        if chunk['type'] in ['subroutine', 'program']
    }
    
    # 解析别名，确保调用关系使用原始名称
    for routine_name, calls in relationships.items():
        for call in calls:
            # 检查函数名是否是别名，如果是则替换为原始名称
            func_name = call['func']
            if func_name in aliases:
                call['func'] = aliases[func_name]
                
    return relationships





